import tensorflow as tf
from tensorflow import keras
import numpy as np
import matplotlib.pyplot as plt
import gzip

# 加载EMNIST数据集
def load_emnist(images_file, labels_file):
    with gzip.open(images_file, 'rb') as f:
        images = np.frombuffer(f.read(), np.uint8, offset=16).reshape(-1, 28, 28)
    with gzip.open(labels_file, 'rb') as f:
        labels = np.frombuffer(f.read(), np.uint8, offset=8)
    return images, labels

# 加载测试数据
test_images, test_labels = load_emnist('emnist-gzip/emnist-byclass-test-images-idx3-ubyte.gz',
                                       'emnist-gzip/emnist-byclass-test-labels-idx1-ubyte.gz')

# 数据预处理
test_images = test_images.reshape((test_images.shape[0], 28, 28, 1)).astype('float32') / 255

# 加载保存的模型
model = keras.models.load_model('emnist_model.h5')

# 评估模型
test_loss, test_acc = model.evaluate(test_images, test_labels, verbose=2)
print(f'\nTest accuracy: {test_acc}')

# 进行预测
predictions = model.predict(test_images)

# EMNIST数据集的标签映射
emnist_labels = [48, 49, 50, 51, 52, 53, 54, 55, 56, 57, 65, 66, 67, 68, 69, 70, 71, 72, 73, 74, 75, 76, 77, 78, 79, 80, 81, 82, 83, 84, 85, 86, 87, 88, 89, 90, 97, 98, 99, 100, 101, 102, 103, 104, 105, 106, 107, 108, 109, 110, 111, 112, 113, 114, 115, 116, 117, 118, 119, 120, 121, 122]

# 显示一些随机选择的预测结果
def plot_image(i, predictions_array, true_label, img):
    true_label, img = true_label[i], img[i]
    plt.grid(False)
    plt.xticks([])
    plt.yticks([])

    plt.imshow(img.reshape((28,28)), cmap=plt.cm.binary)

    predicted_label = np.argmax(predictions_array)
    if predicted_label == true_label:
        color = 'blue'
    else:
        color = 'red'

    plt.xlabel("{} {:2.0f}% ({})".format(chr(emnist_labels[predicted_label]),
                                100*np.max(predictions_array),
                                chr(emnist_labels[true_label])),
                                color=color)

def plot_value_array(i, predictions_array, true_label):
    true_label = true_label[i]
    plt.grid(False)
    plt.xticks(range(62))
    plt.yticks([])
    thisplot = plt.bar(range(62), predictions_array, color="#777777")
    plt.ylim([0, 1])
    predicted_label = np.argmax(predictions_array)

    thisplot[predicted_label].set_color('red')
    thisplot[true_label].set_color('blue')

# 随机选择25个图像进行预测
num_rows = 5
num_cols = 5
num_images = num_rows*num_cols
plt.figure(figsize=(2*2*num_cols, 2*num_rows))
for i in range(num_images):
    idx = np.random.randint(0, len(test_images))
    plt.subplot(num_rows, 2*num_cols, 2*i+1)
    plot_image(idx, predictions[idx], test_labels, test_images)
    plt.subplot(num_rows, 2*num_cols, 2*i+2)
    plot_value_array(idx, predictions[idx], test_labels)
plt.tight_layout()
plt.show()